package cn.demoncat.util.websocket.util;

import cn.demoncat.util.lang.LongUtil;
import cn.demoncat.util.lang.MapUtil;
import cn.demoncat.util.lang.NumberUtil;
import cn.demoncat.util.lang.constant.LongConstant;
import cn.demoncat.util.web.entity.TokenPrincipal;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageHeaders;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.messaging.simp.SimpMessageType;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * WebSocket工具（STOMP协议）
 * 
 * @author 延晓磊
 *
 * @since 2020年3月4日
 */
@Component
public class WebSocketStompUtil {

	@Autowired
	private SimpMessagingTemplate simpMessagingTemplate;
	
	/**
	 * 客户端Session集：{sessionId : user}
	 */
	private static final Map<String, TokenPrincipal> sessions = new ConcurrentHashMap<>();
	/**
	 * 客户端User集：{userId : Set[sessionId]}
	 */
	private static final Map<Long, Set<String>> users = new ConcurrentHashMap<>();
	
	/**
	 * 客户端Session订阅集：{sessionId : Set[queue]}
	 */
	private static final Map<String, Set<String>> sessionQueues = new ConcurrentHashMap<>();
	/**
	 * 客户端User订阅集：{userId : Set[queue]}
	 */
	private static final Map<Long, Set<String>> userQueues = new ConcurrentHashMap<>();
	
	/**
	 * 发送广播消息
	 * 
	 * @param topic	主题：/topic{/topic}
	 * @param msg	消息
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public void sendTopic(String topic, String msg) {
		// 组装主题
		topic = getTopicPath(topic);
		// 广播消息：发送到订阅"/topic{/topic}"的客户端
		simpMessagingTemplate.convertAndSend(topic, msg);
	}
	
	/**
	 * 发送点对点消息 - 到用户
	 * 
	 * @param userId	用户
	 * @param queue		主题：/queue{/queue}
	 * @param msg		消息
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public void sendUser(Long userId, String queue, String msg) {
		// 组装主题
		queue = getQueuePath(queue);
		// 判断user订阅queue
		if (containsUserQueue(userId, queue)) {
			// 点对点消息：发送到订阅"/user/queue{/queue}"的指定用户
			simpMessagingTemplate.convertAndSendToUser(userId.toString(), queue, msg);
		}
	}
	
	/**
	 * 发送点对点消息 - 到用户
	 * 
	 * @param user		用户ID，Principal.name
	 * @param queue		主题：/queue{/queue}
	 * @param msg		消息
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public void sendUser(String user, String queue, String msg) {
		// 组装主题
		queue = getQueuePath(queue);
		// 判断user订阅queue
		if (containsUserQueue(Long.valueOf(user), queue)) {
			// 点对点消息：发送到订阅"/user/queue{/queue}"的指定用户
			simpMessagingTemplate.convertAndSendToUser(user, queue, msg);
		}
	}
	
	/**
	 * 发送点对点消息 - 到会话
	 * 
	 * @param sessionId	会话ID
	 * @param queue	主题：/queue{/queue}
	 * @param msg	消息
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public void sendSession(String sessionId, String queue, String msg) {
		// 组装主题
		queue = getQueuePath(queue);
		// 判断session订阅queue
		if (containsSessionQueue(sessionId, queue)) {
			// 点对点消息：发送到订阅"/user/queue{/queue}"的指定客户端Session
			simpMessagingTemplate.convertAndSendToUser(sessionId, queue, msg, createHeaders(sessionId));
		}
	}
	
	/**
	 * 获取WS-STOMP头信息
	 * 
	 * @param message
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static StompHeaderAccessor getAccess(Message<?> message) {
		return MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
	}
	
	/**
	 * 获取用户Header
	 * 
	 * @param access 	控制器参数，读取WS-STOMP头信息
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static String getHeader(SimpMessageHeaderAccessor access, String header) {
		return access.getFirstNativeHeader(header);
	}
	
	/**
	 * 获取SessionId
	 * 
	 * @param access 	控制器参数，读取WS-STOMP头信息
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static String getSessionId(SimpMessageHeaderAccessor access) {
		return access.getSessionId();
	}
	
	/**
	 * 获取UserId
	 * 
	 * @param access 	控制器参数，读取WS-STOMP头信息
	 * 
	 * @return 如果是登录Token连接，返回用户ID；0表示未登录用户
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static Long getUserId(SimpMessageHeaderAccessor access) {
		if (access.getUser() == null) {
			return LongConstant.N0;
		}
		return LongUtil.valueOf(access.getUser().getName());
	}
	
	/**
	 * 获取User
	 * 
	 * @param access 	控制器参数，读取WS-STOMP头信息
	 * 
	 * @return 如果是登录Token连接，返回用户Principal
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static TokenPrincipal getUser(SimpMessageHeaderAccessor access) {
		if (access.getUser() == null) {
			return null;
		}
		return (TokenPrincipal) access.getUser();
	}
	
	/**
	 * 获取Session用户
	 * 
	 * @param sessionId
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static TokenPrincipal getSessionUser(String sessionId) {
		return sessions.get(sessionId);
	}
	
	/**
	 * 判断是否存在Session
	 * 
	 * @param sessionId
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static boolean containsSession(String sessionId) {
		return sessions.containsKey(sessionId);
	}
	
	/**
	 * 获取Session总数
	 * 
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static int countSession() {
		return sessions.size();
	}
	
	/**
	 * 获取Session集
	 * 
	 * @return {sessionId : user}
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static Map<String, TokenPrincipal> getSessions() {
		return sessions;
	}
	
	/**
	 * 获取用户Session集
	 * 
	 * @param userId
	 * @return sessions
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static Set<String> getUserSessions(Long userId) {
		return users.get(userId);
	}
	
	/**
	 * 判断是否存在用户
	 * 
	 * @param userId
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static boolean containsUser(Long userId) {
		return users.containsKey(userId);
	}
	
	/**
	 * 获取User总数
	 * 
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static int countUser() {
		return users.size();
	}
	
	/**
	 * 获取User集
	 * 
	 * @return {userId : Set[sessionId]}
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static Map<Long, Set<String>> getUsers() {
		return users;
	}
	
	/**
	 * 获取Session订阅集
	 * 
	 * @param sessionId
	 * @return queues
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static Set<String> getSessionQueues(String sessionId) {
		return sessionQueues.get(sessionId);
	}
	
	/**
	 * 判断是否存在Session订阅
	 * 
	 * @param sessionId
	 * @param queue
	 * @return 
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static boolean containsSessionQueue(String sessionId, String queue) {
		return sessionQueues.containsKey(sessionId) && sessionQueues.get(sessionId).contains(queue);
	}
	
	/**
	 * 获取订阅主题的Session
	 * 
	 * @param queue
	 * @return sessions
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static Set<String> getQueueSessions(String queue) {
		// 过滤val包含queue的key
		return sessionQueues.entrySet().parallelStream().filter( entry -> entry.getValue().contains(queue)).map(Map.Entry::getKey).collect(Collectors.toSet());
	}
	
	/**
	 * 获取User订阅集
	 * 
	 * @param userId
	 * @return queues
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static Set<String> getUserQueues(Long userId) {
		return userQueues.get(userId);
	}
	
	/**
	 * 判断是否存在User订阅
	 * 
	 * @param userId
	 * @param queue
	 * @return 
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static boolean containsUserQueue(Long userId, String queue) {
		return userQueues.containsKey(userId) && userQueues.get(userId).contains(queue);
	}
	
	/**
	 * 获取订阅主题的User
	 * 
	 * @param queue
	 * @return users
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static Set<Long> getQueueUsers(String queue) {
		// 过滤val包含queue的key
		return userQueues.entrySet().parallelStream().filter( entry -> entry.getValue().contains(queue)).map(Map.Entry::getKey).collect(Collectors.toSet());
	}
	
	/**
	 * 添加Session
	 * 
	 * 场景：客户端连接成功
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static void addSession(TokenPrincipal user) {
		// Session
		sessions.put(user.getSessionId(), user);
		// User
		if (!NumberUtil.isBlank(user.getUserId())) {
			MapUtil.putSet(users, user.getUserId(), user.getSessionId());
		}
	}
	
	/**
	 * 移除Session
	 * 
	 * 场景：客户端断开连接
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static void removeSession(TokenPrincipal user) {
		// Session
		sessions.remove(user.getSessionId());
		// - Queue
		sessionQueues.remove(user.getSessionId());
		// User
		if (!NumberUtil.isBlank(user.getUserId())) {
			MapUtil.removeSet(users, user.getUserId(), user.getSessionId());
			// - Queue
			userQueues.remove(user.getUserId());
		}
	}
	
	/**
	 * 添加订阅
	 * 
	 * 场景：客户端订阅用户主题
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static void addQueue(TokenPrincipal user, String queue) {
		// 修改主题
		queue = trimPath(queue);
		// Session Queue
		MapUtil.putSet(sessionQueues, user.getSessionId(), queue);
		// User Queue
		if (!NumberUtil.isBlank(user.getUserId())) {
			MapUtil.putSet(userQueues, user.getUserId(), queue);
		}
	}
	
	/**
	 * 移除订阅
	 * 
	 * 场景：客户端取消订阅用户主题
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static void removeQueue(TokenPrincipal user, String queue) {
		// 修改主题
		queue = trimPath(queue);
		// Session Queue
		MapUtil.removeSet(sessionQueues, user.getSessionId(), queue);
		// User Queue
		if (!NumberUtil.isBlank(user.getUserId())) {
			MapUtil.removeSet(userQueues, user.getUserId(), queue);
		}
	}
	
	/**
	 * 获取广播路径
	 * 
	 * @param topic	主题路径
	 * 
	 * @return /topic{/topic}
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static String getTopicPath(String topic) {
		return "/topic" + topic;
	}
	
	/**
	 * 获取点对点路径
	 * 
	 * @param queue	主题路径
	 * 
	 * @return /queue{/queue}
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static String getQueuePath(String queue) {
		return "/queue" + queue;
	}
	
	/**
	 * 修剪客户端订阅主题
	 * 
	 * @param queue	客户端订阅路径：/user/queue/xxx
	 * 
	 * @return /queue/xxx
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static String trimPath(String queue) {
		if (queue.startsWith("/user")) {
			return queue.substring(5);
		}
		return queue;
	}
	
	/**
	 * 是否为客户端订阅的用户主题
	 * 
	 * @param queue	客户端订阅路径：/user/queue/xxx
	 * 
	 * @return 
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	public static boolean isUserPath(String queue) {
		return queue != null && queue.startsWith("/user");
	}
	
	/**
	 * Session发送的请求头
	 * 
	 * @see org.springframework.messaging.simp.annotation.SendToUser
	 * @see org.springframework.messaging.simp.annotation.support.SendToMethodReturnValueHandler#handleReturnValue
	 * 
	 * @param sessionId
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年3月4日
	 */
	private MessageHeaders createHeaders(String sessionId) {
		// 创建消息头
		SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
		// 初始化
		if (simpMessagingTemplate.getHeaderInitializer() != null) {
			simpMessagingTemplate.getHeaderInitializer().initHeaders(headerAccessor);
		}
		// 设置SessionId
		headerAccessor.setSessionId(sessionId);
		// 修改修改
		headerAccessor.setLeaveMutable(true);
		return headerAccessor.getMessageHeaders();
	}
}
